package cc.jinglupeng.kom.data;

import cc.jinglupeng.kom.bean.ZkBroker;
import cc.jinglupeng.kom.bean.ZkPartitionState;
import cc.jinglupeng.kom.pool.ZookeeperPool;
import cc.jinglupeng.kom.util.ZkUtils;
import org.apache.zookeeper.ZooKeeper;
import org.springframework.util.Assert;

import java.util.List;

/**
 * Created by jinglupeng on 4/10/16.
 */
public class ZkData {

    /**
     * 获取所有的 broker
     *
     * @return
     * @throws Exception
     */
    public static ZkBroker[] getAllBrokers() throws Exception {
        List<String> ids = ZkUtils.getChildren("/brokers/ids");
        ZkBroker[] zkBrokers = new ZkBroker[ids.size()];
        int i = 0;
        for (String id : ids) {
            String path = String.format("/brokers/ids/%s", id);
            ZkBroker zkBroker = ZkUtils.getData(path, ZkBroker.class);
            zkBroker.setId(Integer.parseInt(id));
            zkBrokers[i++] = zkBroker;
        }
        return zkBrokers;
    }


    /**
     * 获取所有的 topic
     *
     * @return
     * @throws Exception
     */
    public static List<String> getAllTopic() throws Exception {
        String path = "/brokers/topics";
        return ZkUtils.getChildren(path);
    }


    /**
     * 获取所有的 consumer
     *
     * @return
     * @throws Exception
     */
    public static List<String> getAllConsumer() throws Exception {
        String path = "/consumers";
        return ZkUtils.getChildren(path);
    }

    /**
     * 获取 consumer 消费的所有 topic
     *
     * @param consumer
     * @return
     * @throws Exception
     */
    public static List<String> getConsumerTopics(String consumer) throws Exception {
        Assert.notNull(consumer, "consumer cannot be null.");
        String path = String.format("/consumers/%s/owners", consumer);
        return ZkUtils.getChildren(path);
    }


    public static Integer[] getTopicPartitions(String topic) throws Exception {
        Assert.notNull(topic, "topic cannot be null.");

        String path = String.format("/brokers/topics/%s/partitions", topic);
        List<String> children = ZkUtils.getChildren(path);
        Integer[] partitions = new Integer[children.size()];
        for (int i = 0; i < partitions.length; i++) {
            partitions[i] = Integer.parseInt(children.get(i));
        }
        return partitions;
    }


    public static ZkPartitionState[] getTopicPartitionStates(String topic) throws Exception {
        Assert.notNull(topic, "topic cannot be null.");

        ZooKeeper zooKeeper = ZookeeperPool.borrowZooKeeper();
        ZkPartitionState[] zkPartitionStates = null;
        try {
            String partitionsPath = String.format("/brokers/topics/%s/partitions", topic);
            List<String> partitions = ZkUtils.getChildren(partitionsPath);
            int length = partitions.size();
            zkPartitionStates = new ZkPartitionState[length];
            for (int i = 0; i < length; i++) {
                String partition = partitions.get(i);
                String path = String.format("/brokers/topics/%s/partitions/%s/state", topic, partition);
                zkPartitionStates[i] = ZkUtils.getData(path, ZkPartitionState.class);
                if (zkPartitionStates[i] != null) {
                    zkPartitionStates[i].setTopic(topic);
                    zkPartitionStates[i].setPartition(Integer.parseInt(partition));
                }
            }
        } finally {
            ZookeeperPool.returnZooKeeper(zooKeeper);
        }
        return zkPartitionStates;
    }


    public static Long getConsumerTopicPartitionOffset(String consumer, String topic, int partition) throws Exception {
        Assert.notNull(consumer, "consumer cannot be null.");
        Assert.notNull(topic, "topic cannot be null.");

        String path = String.format("/consumers/%s/offsets/%s/%d", consumer, topic, partition);
        String value = ZkUtils.getData(path);
        if (value == null) return null;
        return Long.parseLong(value);
    }


    public static ZkPartitionState getPartitionState(String topic, int partition) throws Exception {
        Assert.notNull(topic, "topic cannot be null.");

        String path = String.format("/brokers/topics/%s/partitions/%d/state", topic, partition);
        ZkPartitionState zkPartitionState = ZkUtils.getData(path, ZkPartitionState.class);
        if (zkPartitionState == null) return null;
        zkPartitionState.setTopic(topic);
        zkPartitionState.setPartition(partition);
        return zkPartitionState;
    }

}
